"""
MCP 客户端连接器
负责管理与 MCP 服务器的连接和通信
"""

import asyncio
import logging
from typing import Dict, List, Optional, Union
from urllib.parse import urlparse

class MCPConnector:
    """MCP 客户端连接器
    
    管理与多个 MCP 服务器的连接,提供统一的请求接口
    支持连接池管理、请求重试和错误处理
    """
    
    def __init__(self, servers: List[str], config: Optional[Dict] = None):
        """初始化 MCP 连接器
        
        Args:
            servers: MCP 服务器 URL 列表
            config: 连接器配置
                max_retries: 最大重试次数
                timeout: 请求超时时间
                pool_size: 连接池大小
        """
        self.servers = servers
        self.config = config or {}
        self.connections = {}
        self.logger = logging.getLogger(__name__)
        
        # 初始化连接池
        self._init_connection_pools()
        
    def _init_connection_pools(self) -> None:
        """初始化服务器连接池"""
        for server_url in self.servers:
            parsed_url = urlparse(server_url)
            server_type = parsed_url.scheme
            
            # 根据服务器类型创建连接池
            if server_type == "mysql":
                from ..servers.mysql_server import MySQLMCPServer
                config = self._parse_mysql_url(server_url)
                self.connections[server_url] = MySQLMCPServer(config)
            elif server_type == "elasticsearch":
                from ..servers.es_server import ElasticsearchMCPServer
                config = self._parse_es_url(server_url)
                self.connections[server_url] = ElasticsearchMCPServer(config)
            elif server_type == "redis":
                from ..servers.redis_server import RedisMCPServer
                config = self._parse_redis_url(server_url)
                self.connections[server_url] = RedisMCPServer(config)
            elif server_type == "kafka":
                from ..servers.kafka_server import KafkaMCPServer
                config = self._parse_kafka_url(server_url)
                self.connections[server_url] = KafkaMCPServer(config)
            else:
                raise ValueError(f"Unsupported server type: {server_type}")
                
    def _parse_mysql_url(self, url: str) -> Dict:
        """解析 MySQL 连接 URL"""
        parsed = urlparse(url)
        return {
            "host": parsed.hostname,
            "port": parsed.port or 3306,
            "user": parsed.username,
            "password": parsed.password,
            "database": parsed.path.strip("/"),
            "max_connections": self.config.get("pool_size", 10)
        }
        
    def _parse_es_url(self, url: str) -> Dict:
        """解析 Elasticsearch 连接 URL"""
        parsed = urlparse(url)
        return {
            "hosts": [f"{parsed.hostname}:{parsed.port or 9200}"],
            "http_auth": (parsed.username, parsed.password) if parsed.username else None,
            "scheme": parsed.scheme,
            "verify_certs": self.config.get("verify_certs", True)
        }
        
    def _parse_redis_url(self, url: str) -> Dict:
        """解析 Redis 连接 URL"""
        parsed = urlparse(url)
        return {
            "host": parsed.hostname,
            "port": parsed.port or 6379,
            "password": parsed.password,
            "db": int(parsed.path.strip("/") or 0),
            "decode_responses": True
        }
        
    def _parse_kafka_url(self, url: str) -> Dict:
        """解析 Kafka 连接 URL"""
        parsed = urlparse(url)
        return {
            "bootstrap_servers": [f"{parsed.hostname}:{parsed.port or 9092}"],
            "security_protocol": "PLAINTEXT",
            "client_id": self.config.get("client_id", "mcp_client")
        }
        
    async def connect(self) -> None:
        """连接所有配置的服务器"""
        for server_url, server in self.connections.items():
            try:
                await server.connect()
                self.logger.info(f"Connected to {server_url}")
            except Exception as e:
                self.logger.error(f"Failed to connect to {server_url}: {e}")
                raise
                
    async def disconnect(self) -> None:
        """断开所有服务器连接"""
        for server_url, server in self.connections.items():
            try:
                await server.disconnect()
                self.logger.info(f"Disconnected from {server_url}")
            except Exception as e:
                self.logger.error(f"Error disconnecting from {server_url}: {e}")
                
    async def send_request(
        self,
        server_url: str,
        request: Dict,
        retry_count: int = 0
    ) -> Dict:
        """发送请求到指定服务器
        
        Args:
            server_url: 服务器 URL
            request: MCP 请求
            retry_count: 当前重试次数
            
        Returns:
            Dict: MCP 响应
            
        Raises:
            ConnectionError: 连接错误
            TimeoutError: 请求超时
        """
        max_retries = self.config.get("max_retries", 3)
        timeout = self.config.get("timeout", 30)
        
        try:
            server = self.connections[server_url]
            
            # 设置超时
            async with asyncio.timeout(timeout):
                response = await server.handle_request(request)
                
            return response
            
        except (ConnectionError, TimeoutError) as e:
            if retry_count < max_retries:
                # 重试
                await asyncio.sleep(2 ** retry_count)  # 指数退避
                return await self.send_request(
                    server_url,
                    request,
                    retry_count + 1
                )
            else:
                self.logger.error(f"Request failed after {max_retries} retries")
                raise
                
    async def broadcast_request(self, request: Dict) -> List[Dict]:
        """向所有服务器广播请求
        
        Args:
            request: MCP 请求
            
        Returns:
            List[Dict]: 所有服务器的响应列表
        """
        tasks = []
        for server_url in self.servers:
            task = asyncio.create_task(
                self.send_request(server_url, request)
            )
            tasks.append(task)
            
        responses = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 过滤出成功的响应
        valid_responses = [
            resp for resp in responses
            if isinstance(resp, dict) and resp.get("status") == "success"
        ]
        
        return valid_responses
        
    async def query_all_sources(self, query: str) -> List[Dict]:
        """查询所有数据源
        
        Args:
            query: 查询语句
            
        Returns:
            List[Dict]: 合并后的查询结果
        """
        request = {
            "action": "query",
            "query": query
        }
        
        responses = await self.broadcast_request(request)
        
        # 合并所有数据源的结果
        all_results = []
        for response in responses:
            if response.get("data"):
                all_results.extend(response["data"])
                
        return all_results 